package com.sheepone.util.weka;

import weka.classifiers.AbstractClassifier;
import weka.classifiers.trees.M5P;
import weka.core.DenseInstance;
import weka.core.Instance;
import weka.core.Instances;
import weka.core.SerializationHelper;

/**
 * @author Miss.杨
 * @description
 * @since 2024/3/31 - 11:05
 */
public class M5PUtil {

    private static Instance convertToIndependentVariable(double value, int index,int numAttributes,Instances data) {
        Instance wekaInstance = new DenseInstance(numAttributes);
        wekaInstance.setValue(index, value); // 设置自变量值
        wekaInstance.setDataset(data);
        return wekaInstance;
    }

    /**
     * 使用M5P模型对给定的二维数组进行了非线性回归拟合
     *
     * @param data 拟合数据
     * @param dependentVariableIndex 因变量的下标
     * @return
     */
    public static M5P build(double[][] data, int dependentVariableIndex) {
        return build(WekaUtil.convertToInstances(data,dependentVariableIndex));
    }

    /**
     * 使用M5P模型对给定的Instances进行了非线性回归拟合
     *
     * @param instances 拟合数据
     * @return 模型
     */
    public static M5P build(Instances instances) {
        M5P m5P = new M5P();
        try {
            // 构建M5P回归树模型
            m5P.buildClassifier(instances);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        return m5P;
    }

    /**
     * 预测自变量在改拟合曲线的结果因变量值
     *
     * @param classifier          拟合曲线
     * @param independentVariable 自变量
     * @param independentVariableIndex 自变量索引
     * @param numAttributes            特征数
     * @param originalData        原始数据
     * @return 预测结果
     */
    public static double predict(AbstractClassifier classifier, double independentVariable, int independentVariableIndex, int numAttributes, double[][] originalData) {
        return predict(classifier, independentVariable, independentVariableIndex,numAttributes, WekaUtil.convertToInstances(originalData));
    }

    /**
     * 预测自变量在改拟合曲线的结果因变量值
     *
     * @param classifier               拟合曲线
     * @param independentVariable      自变量
     * @param independentVariableIndex 自变量索引
     * @param numAttributes            特征数
     * @param originData               原始数据
     * @return 预测结果
     */
    public static double predict(AbstractClassifier classifier, double independentVariable,int independentVariableIndex, int numAttributes,Instances originData) {
        return predict(classifier, convertToIndependentVariable(independentVariable,independentVariableIndex,numAttributes,originData));
    }


    /**
     * 预测自变量在改拟合曲线的结果因变量值
     *
     * @param classifier          拟合曲线
     * @param independentVariable 自变量
     * @return 预测结果
     */
    private static double predict(AbstractClassifier classifier, Instance independentVariable) {
        try {
            return classifier.classifyInstance(independentVariable);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 保存M5P模型
     *
     * @param m5p 模型
     * @param filePath 文件路径
     */

    public static void save(M5P m5p, String filePath)  {
        try {
            SerializationHelper.write(filePath, m5p);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 加载M5P模型
     * @param filePath 文件路径
     * @return 模型
     */
    public static M5P load(String filePath)  {
        try {
            return (M5P) SerializationHelper.read(filePath);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }
}
